{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------\n",
      "data-single processing\n",
      "process_med\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/sjy1203/anaconda3/lib/python3.6/site-packages/ipykernel_launcher.py:278: DtypeWarning: Columns (11) have mixed types. Specify dtype option on import or set low_memory=False.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "process_diag\n",
      "filter diag\n",
      "filter_patient\n",
      "--------------------\n",
      "data-multi processing \n",
      "--------------------\n",
      "data-single stat\n",
      "#patients  (30745,)\n",
      "#clinical events  30745\n",
      "#diagnosis  1997\n",
      "#med  323\n",
      "#avg of diagnoses  10.820458611156285\n",
      "#avg of medicines  13.759277931370955\n",
      "#avg of vists  1.0\n",
      "#max of diagnoses  39\n",
      "#max of medicines  52\n",
      "#max of visit  1\n",
      "--------------------\n",
      "data_multi stat\n",
      "#patients  (6350,)\n",
      "#clinical events  15016\n",
      "#diagnosis  1958\n",
      "#med  145\n",
      "#avg of diagnoses  10.514717634523175\n",
      "#avg of medicines  8.80420884389984\n",
      "#avg of vists  2.3647244094488187\n",
      "#max of diagnoses  128\n",
      "#max of medicines  55\n",
      "#max of visit  29\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>SUBJECT_ID</th>\n",
       "      <th>HADM_ID</th>\n",
       "      <th>ICD9_CODE</th>\n",
       "      <th>ATC4</th>\n",
       "      <th>PRO_CODE</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>17</td>\n",
       "      <td>161087</td>\n",
       "      <td>[4239, 5119, 78551, 4589, 311, 7220, 71946, 2724]</td>\n",
       "      <td>[N02B, A01A, A02B, A06A, B05C, A12A, A12C, C01...</td>\n",
       "      <td>[3731, 8872, 3893]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>17</td>\n",
       "      <td>194023</td>\n",
       "      <td>[7455, 45829, V1259, 2724]</td>\n",
       "      <td>[N02B, A01A, A02B, A06A, A12A, B05C, A12C, C01...</td>\n",
       "      <td>[3571, 3961, 8872]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>21</td>\n",
       "      <td>109451</td>\n",
       "      <td>[41071, 78551, 5781, 5849, 40391, 4280, 4592, ...</td>\n",
       "      <td>[A06A, C07A, A12A, A02A, J01M, C02A, B05C, B01...</td>\n",
       "      <td>[0066, 3761, 3950, 3606, 0042, 0047, 3895, 399...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>21</td>\n",
       "      <td>111970</td>\n",
       "      <td>[0388, 78552, 40391, 42731, 70709, 5119, 6823,...</td>\n",
       "      <td>[A06A, B05C, A12C, A07A, N02B, B01A, N06A, A01...</td>\n",
       "      <td>[3995, 8961, 0014]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>23</td>\n",
       "      <td>124321</td>\n",
       "      <td>[2252, 3485, 78039, 4241, 4019, 2720, 2724, V4...</td>\n",
       "      <td>[C07A, N02B, A02B, H03A, N03A, A01A, N05A, C09...</td>\n",
       "      <td>[0151]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>23</td>\n",
       "      <td>152223</td>\n",
       "      <td>[41401, 4111, 4241, V4582, 2724, 4019, 60000, ...</td>\n",
       "      <td>[N02B, A01A, A02B, A06A, A12A, B05C, A12C, C01...</td>\n",
       "      <td>[3613, 3615, 3961, 8872, 9904, 9905, 9907]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>34</td>\n",
       "      <td>115799</td>\n",
       "      <td>[41071, 4280, 4254, 42731, 9971, 4260, 41401, ...</td>\n",
       "      <td>[C07A, N02B, B01A, B05C, C10A, D11A, C03C, H03...</td>\n",
       "      <td>[3723, 8852, 8856]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>34</td>\n",
       "      <td>144319</td>\n",
       "      <td>[42789, 42822, 4263, 41401, V5861, 4280, 2449,...</td>\n",
       "      <td>[A01A, H03A, D11A, N05C, A04A, C10A]</td>\n",
       "      <td>[3772, 3783, 8945]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>36</td>\n",
       "      <td>122659</td>\n",
       "      <td>[99831, 99811, 41511, 4538, 99672, 496, 41401,...</td>\n",
       "      <td>[A01A, A12A, B05C, C01C, A07A, N01A, C07A, B01...</td>\n",
       "      <td>[3479, 7761, 3932, 3403, 8674, 7851, 3893, 9904]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>36</td>\n",
       "      <td>165660</td>\n",
       "      <td>[55321, 41511, 5185, 486, 99739, 5180, 41400, ...</td>\n",
       "      <td>[A01A, A12A, A12C, N02B, B01A, C09A, C10A, C08...</td>\n",
       "      <td>[5361, 3479, 7861, 544, 9977, 9671, 9604]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   SUBJECT_ID  HADM_ID                                          ICD9_CODE  \\\n",
       "0          17   161087  [4239, 5119, 78551, 4589, 311, 7220, 71946, 2724]   \n",
       "1          17   194023                         [7455, 45829, V1259, 2724]   \n",
       "2          21   109451  [41071, 78551, 5781, 5849, 40391, 4280, 4592, ...   \n",
       "3          21   111970  [0388, 78552, 40391, 42731, 70709, 5119, 6823,...   \n",
       "4          23   124321  [2252, 3485, 78039, 4241, 4019, 2720, 2724, V4...   \n",
       "5          23   152223  [41401, 4111, 4241, V4582, 2724, 4019, 60000, ...   \n",
       "6          34   115799  [41071, 4280, 4254, 42731, 9971, 4260, 41401, ...   \n",
       "7          34   144319  [42789, 42822, 4263, 41401, V5861, 4280, 2449,...   \n",
       "8          36   122659  [99831, 99811, 41511, 4538, 99672, 496, 41401,...   \n",
       "9          36   165660  [55321, 41511, 5185, 486, 99739, 5180, 41400, ...   \n",
       "\n",
       "                                                ATC4  \\\n",
       "0  [N02B, A01A, A02B, A06A, B05C, A12A, A12C, C01...   \n",
       "1  [N02B, A01A, A02B, A06A, A12A, B05C, A12C, C01...   \n",
       "2  [A06A, C07A, A12A, A02A, J01M, C02A, B05C, B01...   \n",
       "3  [A06A, B05C, A12C, A07A, N02B, B01A, N06A, A01...   \n",
       "4  [C07A, N02B, A02B, H03A, N03A, A01A, N05A, C09...   \n",
       "5  [N02B, A01A, A02B, A06A, A12A, B05C, A12C, C01...   \n",
       "6  [C07A, N02B, B01A, B05C, C10A, D11A, C03C, H03...   \n",
       "7               [A01A, H03A, D11A, N05C, A04A, C10A]   \n",
       "8  [A01A, A12A, B05C, C01C, A07A, N01A, C07A, B01...   \n",
       "9  [A01A, A12A, A12C, N02B, B01A, C09A, C10A, C08...   \n",
       "\n",
       "                                            PRO_CODE  \n",
       "0                                 [3731, 8872, 3893]  \n",
       "1                                 [3571, 3961, 8872]  \n",
       "2  [0066, 3761, 3950, 3606, 0042, 0047, 3895, 399...  \n",
       "3                                 [3995, 8961, 0014]  \n",
       "4                                             [0151]  \n",
       "5         [3613, 3615, 3961, 8872, 9904, 9905, 9907]  \n",
       "6                                 [3723, 8852, 8856]  \n",
       "7                                 [3772, 3783, 8945]  \n",
       "8   [3479, 7761, 3932, 3403, 8674, 7851, 3893, 9904]  \n",
       "9          [5361, 3479, 7861, 544, 9977, 9671, 9604]  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "\n",
    "med_file = 'PRESCRIPTIONS.csv'\n",
    "diag_file = 'DIAGNOSES_ICD.csv'\n",
    "ndc2atc_file = 'ndc2atc_level4.csv'\n",
    "ddi_file = 'drug-DDI.csv'\n",
    "cid_atc = 'drug-atc.csv'\n",
    "patient_info_file='./gather_firstday.csv'\n",
    "\n",
    "\n",
    "def process_med():\n",
    "    print('process_med')\n",
    "    med_pd = pd.read_csv(med_file, dtype={'NDC':'category'})\n",
    "    # filter\n",
    "    med_pd.drop(columns=['ROW_ID','DRUG_TYPE','DRUG_NAME_POE','DRUG_NAME_GENERIC',\n",
    "                     'FORMULARY_DRUG_CD','GSN','PROD_STRENGTH','DOSE_VAL_RX',\n",
    "                     'DOSE_UNIT_RX','FORM_VAL_DISP','FORM_UNIT_DISP','FORM_UNIT_DISP',\n",
    "                      'ROUTE','ENDDATE','DRUG'], axis=1, inplace=True)\n",
    "    med_pd.drop(index = med_pd[med_pd['NDC'] == '0'].index, axis=0, inplace=True)\n",
    "    med_pd.fillna(method='pad', inplace=True)\n",
    "    med_pd.dropna(inplace=True)\n",
    "    med_pd.drop_duplicates(inplace=True)\n",
    "    med_pd['ICUSTAY_ID'] = med_pd['ICUSTAY_ID'].astype('int64')\n",
    "    med_pd['STARTDATE'] = pd.to_datetime(med_pd['STARTDATE'], format='%Y-%m-%d %H:%M:%S')    \n",
    "    med_pd.sort_values(by=['SUBJECT_ID', 'HADM_ID', 'ICUSTAY_ID', 'STARTDATE'], inplace=True)\n",
    "    med_pd = med_pd.reset_index(drop=True)\n",
    "    \n",
    "    def filter_first24hour_med(med_pd):\n",
    "        med_pd_new = med_pd.drop(columns=['NDC'])\n",
    "        med_pd_new = med_pd_new.groupby(by=['SUBJECT_ID','HADM_ID','ICUSTAY_ID']).head([1]).reset_index(drop=True)\n",
    "        med_pd_new = pd.merge(med_pd_new, med_pd, on=['SUBJECT_ID','HADM_ID','ICUSTAY_ID','STARTDATE'])\n",
    "        med_pd_new = med_pd_new.drop(columns=['STARTDATE'])\n",
    "        return med_pd_new\n",
    "    med_pd = filter_first24hour_med(med_pd) # or next line\n",
    "#     med_pd = med_pd.drop(columns=['STARTDATE'])\n",
    "    \n",
    "    med_pd = med_pd.drop(columns=['ICUSTAY_ID'])\n",
    "    med_pd = med_pd.drop_duplicates()\n",
    "    \n",
    "    return med_pd.reset_index(drop=True)\n",
    "\n",
    "def process_diag():\n",
    "    print('process_diag')\n",
    "    \n",
    "    diag_pd = pd.read_csv(diag_file)\n",
    "    diag_pd.dropna(inplace=True)\n",
    "    diag_pd.drop(columns=['SEQ_NUM','ROW_ID'],inplace=True)\n",
    "    diag_pd.drop_duplicates(inplace=True)\n",
    "    diag_pd.sort_values(by=['SUBJECT_ID','HADM_ID'], inplace=True)\n",
    "    return diag_pd.reset_index(drop=True)\n",
    "\n",
    "def process_side():\n",
    "    print('process_side')\n",
    "    \n",
    "    side_pd = pd.read_csv(patient_info_file)\n",
    "    # just use demographic information to avoid future information leak such as lab test and lab measurements\n",
    "    side_pd = side_pd[['subject_id', 'hadm_id', 'icustay_id',\n",
    "                       'gender_male', 'admission_type', 'first_icu_stay', 'admission_age', \n",
    "                       'ethnicity', 'weight', 'height']]\n",
    "    \n",
    "    #process side_information\n",
    "    side_pd = side_pd.dropna(thresh=4)\n",
    "    side_pd.fillna(side_pd.mean(), inplace=True)\n",
    "    side_pd = side_pd.groupby(by=['subject_id', 'hadm_id']).head([1]).reset_index(drop=True)\n",
    "    side_pd = pd.concat([side_pd, pd.get_dummies(side_pd['ethnicity'])],axis=1)\n",
    "    side_pd.drop(columns=['ethnicity', 'icustay_id'], inplace=True)\n",
    "    side_pd.rename(columns={'subject_id':'SUBJECT_ID', 'hadm_id':'HADM_ID'}, inplace=True)\n",
    "    return side_pd.reset_index(drop=True)\n",
    "\n",
    "def ndc2atc4(med_pd):\n",
    "    with open('ndc2rxnorm_mapping.txt', 'r') as f:\n",
    "        ndc2rxnorm = eval(f.read())\n",
    "    med_pd['RXCUI'] = med_pd['NDC'].map(ndc2rxnorm)\n",
    "    med_pd.dropna(inplace=True)\n",
    "\n",
    "    rxnorm2atc = pd.read_csv('ndc2atc_level4.csv')\n",
    "    rxnorm2atc = rxnorm2atc.drop(columns=['YEAR','MONTH','NDC'])\n",
    "    rxnorm2atc.drop_duplicates(subset=['RXCUI'], inplace=True)\n",
    "    med_pd.drop(index = med_pd[med_pd['RXCUI'].isin([''])].index, axis=0, inplace=True)\n",
    "    \n",
    "    med_pd['RXCUI'] = med_pd['RXCUI'].astype('int64')\n",
    "    med_pd = med_pd.reset_index(drop=True)\n",
    "    med_pd = med_pd.merge(rxnorm2atc, on=['RXCUI'])\n",
    "    med_pd.drop(columns=['NDC', 'RXCUI'], inplace=True)\n",
    "#     med_pd = med_pd.rename(columns={'ATC4':'NDC'})\n",
    "    med_pd['ATC4'] = med_pd['ATC4'].map(lambda x: x[:5])\n",
    "    med_pd = med_pd.drop_duplicates()    \n",
    "    med_pd = med_pd.reset_index(drop=True)\n",
    "    return med_pd\n",
    "\n",
    "def filter_pro(pro_pd):\n",
    "    pro_count = pro_pd.groupby(by=['ICD9_CODE']).size().reset_index().rename(columns={0:'count'}).sort_values(by=['count'],ascending=False).reset_index(drop=True)\n",
    "    pro_pd = pro_pd[pro_pd['ICD9_CODE'].isin(pro_count.loc[:1000, 'ICD9_CODE'])]\n",
    "    \n",
    "    return pro_pd.reset_index(drop=True)    \n",
    "\n",
    "def filter_diag(diag_pd, num=128):\n",
    "    print('filter diag')\n",
    "    diag_count = diag_pd.groupby(by=['ICD9_CODE']).size().reset_index().rename(columns={0:'count'}).sort_values(by=['count'],ascending=False).reset_index(drop=True)\n",
    "    diag_pd = diag_pd[diag_pd['ICD9_CODE'].isin(diag_count.loc[:num, 'ICD9_CODE'])]\n",
    "    \n",
    "    return diag_pd.reset_index(drop=True)\n",
    "\n",
    "def filter_med(med_pd):\n",
    "    med_count = med_pd.groupby(by=['ATC4']).size().reset_index().rename(columns={0:'count'}).sort_values(by=['count'],ascending=False).reset_index(drop=True)\n",
    "    med_pd = med_pd[med_pd['ATC4'].isin(med_count.loc[:299, 'ATC4'])]\n",
    "    \n",
    "    return med_pd.reset_index(drop=True)\n",
    "\n",
    "# visit filter\n",
    "def filter_by_visit_range(data_pd, v_range=(1, 2)):\n",
    "    a = data_pd[['SUBJECT_ID', 'HADM_ID']].groupby(by='SUBJECT_ID')['HADM_ID'].unique().reset_index()\n",
    "    a['HADM_ID_Len'] = a['HADM_ID'].map(lambda x:len(x))\n",
    "    a = a[(a['HADM_ID_Len'] >= v_range[0]) & (a['HADM_ID_Len'] < v_range[1])] \n",
    "    data_pd_filter = a.reset_index(drop=True)    \n",
    "    data_pd = data_pd.merge(data_pd_filter[['SUBJECT_ID']], on='SUBJECT_ID', how='inner')\n",
    "    return data_pd.reset_index(drop=True)\n",
    "\n",
    "def process_all(visit_range=(1,2)):\n",
    "    # get med and diag (visit>=2)\n",
    "    med_pd = process_med()\n",
    "    med_pd = ndc2atc4(med_pd)\n",
    "#     med_pd = filter_300_most_med(med_pd)\n",
    "    med_pd = filter_by_visit_range(med_pd, visit_range)\n",
    "    \n",
    "    diag_pd = process_diag()\n",
    "    diag_pd = filter_diag(diag_pd, num=1999)\n",
    "\n",
    "#     side_pd = process_side()\n",
    "    \n",
    "#     pro_pd = process_procedure()\n",
    "#     pro_pd = filter_1000_most_pro(pro_pd)\n",
    "    \n",
    "    med_pd_key = med_pd[['SUBJECT_ID', 'HADM_ID']].drop_duplicates()\n",
    "    diag_pd_key = diag_pd[['SUBJECT_ID', 'HADM_ID']].drop_duplicates()\n",
    "#     pro_pd_key = pro_pd[['SUBJECT_ID', 'HADM_ID']].drop_duplicates()\n",
    "#     side_pd_key = side_pd[['SUBJECT_ID', 'HADM_ID']].drop_duplicates()\n",
    "    \n",
    "    combined_key = med_pd_key.merge(diag_pd_key, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "#     combined_key = combined_key.merge(pro_pd_key, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "#     combined_key = combined_key.merge(side_pd_key, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "    \n",
    "    diag_pd = diag_pd.merge(combined_key, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "    med_pd = med_pd.merge(combined_key, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "#     side_pd = side_pd.merge(combined_key, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "#     pro_pd = pro_pd.merge(combined_key, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "    \n",
    "    # flatten and merge\n",
    "    diag_pd = diag_pd.groupby(by=['SUBJECT_ID','HADM_ID'])['ICD9_CODE'].unique().reset_index()  \n",
    "    med_pd = med_pd.groupby(by=['SUBJECT_ID', 'HADM_ID'])['ATC4'].unique().reset_index()\n",
    "#     pro_pd = pro_pd.groupby(by=['SUBJECT_ID','HADM_ID'])['ICD9_CODE'].unique().reset_index().rename(columns={'ICD9_CODE':'PRO_CODE'})  \n",
    "    diag_pd['ICD9_CODE'] = diag_pd['ICD9_CODE'].map(lambda x: list(x))\n",
    "    med_pd['ATC4'] = med_pd['ATC4'].map(lambda x: list(x))\n",
    "#     pro_pd['PRO_CODE'] = pro_pd['PRO_CODE'].map(lambda x: list(x))\n",
    "    data = diag_pd.merge(med_pd, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "#     data = data.merge(side_pd, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "#     data = data.merge(pro_pd, on=['SUBJECT_ID', 'HADM_ID'], how='inner')\n",
    "#     data['ICD9_CODE_Len'] = data['ICD9_CODE'].map(lambda x: len(x))\n",
    "#     data['NDC_Len'] = data['NDC'].map(lambda x: len(x))\n",
    "    return data\n",
    "\n",
    "def filter_patient(data, dx_range=(2, np.inf), rx_range=(2, np.inf)):\n",
    "    print('filter_patient')\n",
    "    \n",
    "    drop_subject_ls = []\n",
    "    for subject_id in data['SUBJECT_ID'].unique():\n",
    "        item_data = data[data['SUBJECT_ID'] == subject_id]\n",
    "        \n",
    "        for index, row in item_data.iterrows():\n",
    "            dx_len = len(list(row['ICD9_CODE']))\n",
    "            rx_len = len(list(row['ATC4']))\n",
    "            if  dx_len < dx_range[0] or dx_len > dx_range[1]:\n",
    "                drop_subject_ls.append(subject_id)\n",
    "                break\n",
    "            if  rx_len < rx_range[0] or rx_len > rx_range[1]:\n",
    "                drop_subject_ls.append(subject_id)\n",
    "                break\n",
    "    data.drop(index = data[data['SUBJECT_ID'].isin(drop_subject_ls)].index, axis=0, inplace=True)\n",
    "    return data.reset_index(drop=True)\n",
    "    \n",
    "def statistics(data):\n",
    "    print('#patients ', data['SUBJECT_ID'].unique().shape)\n",
    "    print('#clinical events ', len(data))\n",
    "    \n",
    "    diag = data['ICD9_CODE'].values\n",
    "    med = data['ATC4'].values\n",
    "    \n",
    "    unique_diag = set([j for i in diag for j in list(i)])\n",
    "    unique_med = set([j for i in med for j in list(i)])\n",
    "    \n",
    "    print('#diagnosis ', len(unique_diag))\n",
    "    print('#med ', len(unique_med))\n",
    "    \n",
    "    avg_diag = 0\n",
    "    avg_med = 0\n",
    "    max_diag = 0\n",
    "    max_med = 0\n",
    "    cnt = 0\n",
    "    max_visit = 0\n",
    "    avg_visit = 0\n",
    "\n",
    "    for subject_id in data['SUBJECT_ID'].unique():\n",
    "        item_data = data[data['SUBJECT_ID'] == subject_id]\n",
    "        x = []\n",
    "        y = []\n",
    "        visit_cnt = 0\n",
    "        for index, row in item_data.iterrows():\n",
    "            visit_cnt += 1\n",
    "            cnt += 1\n",
    "            x.extend(list(row['ICD9_CODE']))\n",
    "            y.extend(list(row['ATC4']))\n",
    "        x = set(x)\n",
    "        y = set(y)\n",
    "        avg_diag += len(x)\n",
    "        avg_med += len(y)\n",
    "        avg_visit += visit_cnt\n",
    "        if len(x) > max_diag:\n",
    "            max_diag = len(x)\n",
    "        if len(y) > max_med:\n",
    "            max_med = len(y) \n",
    "        if visit_cnt > max_visit:\n",
    "            max_visit = visit_cnt\n",
    "        \n",
    "\n",
    "        \n",
    "    print('#avg of diagnoses ', avg_diag/ cnt)\n",
    "    print('#avg of medicines ', avg_med/ cnt)\n",
    "    print('#avg of vists ', avg_visit/ len(data['SUBJECT_ID'].unique()))\n",
    "    \n",
    "\n",
    "    print('#max of diagnoses ', max_diag)\n",
    "    print('#max of medicines ', max_med)\n",
    "    print('#max of visit ', max_visit)\n",
    "\n",
    "def run(visit_range=(1,2)):\n",
    "    data = process_all(visit_range)\n",
    "    data = filter_patient(data)\n",
    "    \n",
    "    # unique code save\n",
    "    diag = data['ICD9_CODE'].values\n",
    "    med = data['ATC4'].values\n",
    "    unique_diag = set([j for i in diag for j in list(i)])\n",
    "    unique_med = set([j for i in med for j in list(i)])\n",
    "    \n",
    "    return data, unique_diag, unique_med\n",
    "\n",
    "def load_gamenet_multi_visit_data(file_name='data_gamenet.pkl'):\n",
    "    data = pd.read_pickle(file_name)\n",
    "    data.rename(columns={'NDC':'ATC4'}, inplace=True)\n",
    "    data.drop(columns=['PRO_CODE', 'NDC_Len'], axis=1, inplace=True)\n",
    "    \n",
    "    # unique code save\n",
    "    diag = data['ICD9_CODE'].values\n",
    "    med = data['ATC4'].values\n",
    "    unique_diag = set([j for i in diag for j in list(i)])\n",
    "    unique_med = set([j for i in med for j in list(i)])\n",
    "    return data, unique_diag, unique_med\n",
    "\n",
    "def load_gamenet_multi_visit_data_with_pro(file_name='data_gamenet.pkl'):\n",
    "    data = pd.read_pickle(file_name)\n",
    "    data.rename(columns={'NDC':'ATC4'}, inplace=True)\n",
    "    data.drop(columns=['NDC_Len'], axis=1, inplace=True)\n",
    "    \n",
    "    # unique code save\n",
    "    diag = data['ICD9_CODE'].values\n",
    "    med = data['ATC4'].values\n",
    "    pro = data['PRO_CODE'].values\n",
    "    unique_diag = set([j for i in diag for j in list(i)])\n",
    "    unique_med = set([j for i in med for j in list(i)])\n",
    "    unique_pro = set([j for i in pro for j in list(i)])\n",
    "    \n",
    "    return data, unique_pro, unique_diag, unique_med\n",
    "    \n",
    "def main():\n",
    "    print('-'*20 + '\\ndata-single processing')\n",
    "    data_single_visit, diag1, med1 = run(visit_range=(1,2))\n",
    "    print('-'*20 + '\\ndata-multi processing ')\n",
    "    data_multi_visit, pro, diag2, med2 = load_gamenet_multi_visit_data_with_pro()\n",
    "#     med_diag_pair = gen_med_diag_pair(data)\n",
    "    \n",
    "    unique_diag = diag1 | diag2\n",
    "    unique_med = med1 | med2\n",
    "    with open('dx-vocab.txt', 'w') as fout:\n",
    "        for code in unique_diag:\n",
    "            fout.write(code + '\\n')\n",
    "    with open('rx-vocab.txt', 'w') as fout:\n",
    "        for code in unique_med:\n",
    "            fout.write(code + '\\n')\n",
    "            \n",
    "    with open('rx-vocab-multi.txt', 'w') as fout:\n",
    "        for code in med2:\n",
    "            fout.write(code + '\\n')\n",
    "    with open('dx-vocab-multi.txt', 'w') as fout:\n",
    "        for code in diag2:\n",
    "            fout.write(code + '\\n')\n",
    "    with open('px-vocab-multi.txt', 'w') as fout:\n",
    "        for code in pro:\n",
    "            fout.write(code + '\\n')\n",
    "            \n",
    "    # save data\n",
    "    data_single_visit.to_pickle('data-single-visit.pkl')\n",
    "    data_multi_visit.to_pickle('data-multi-visit.pkl')\n",
    "    \n",
    "#     med_diag_pair.to_pickle('med_diag.pkl')\n",
    "#     print('med2diag len:', len(med_diag_pair))\n",
    "    \n",
    "    print('-'*20 + '\\ndata-single stat')\n",
    "    statistics(data_single_visit)\n",
    "    print('-'*20 + '\\ndata_multi stat')\n",
    "    statistics(data_multi_visit)\n",
    "    \n",
    "    return data_single_visit, data_multi_visit\n",
    "\n",
    "data_single_visit, data_multi_visit = main()\n",
    "data_multi_visit.head(10)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train size: 4233, eval size: 1058, test size: 1059\n"
     ]
    }
   ],
   "source": [
    "# split train, eval and test dataset\n",
    "import random\n",
    "from random import shuffle\n",
    "random.seed(1203)\n",
    "\n",
    "def split_dataset(data_path='data-multi-visit.pkl'):\n",
    "    data = pd.read_pickle(data_path)\n",
    "    sample_id = data['SUBJECT_ID'].unique()\n",
    "    \n",
    "    random_number = [i for i in range(len(sample_id))]\n",
    "#     shuffle(random_number)\n",
    "    \n",
    "    train_id = sample_id[random_number[:int(len(sample_id)*2/3)]]\n",
    "    eval_id = sample_id[random_number[int(len(sample_id)*2/3): int(len(sample_id)*5/6)]]\n",
    "    test_id = sample_id[random_number[int(len(sample_id)*5/6):]]\n",
    "    \n",
    "    def ls2file(list_data, file_name):\n",
    "        with open(file_name, 'w') as fout:\n",
    "            for item in list_data:\n",
    "                fout.write(str(item) + '\\n')\n",
    "    \n",
    "    ls2file(train_id, 'train-id.txt')\n",
    "    ls2file(eval_id, 'eval-id.txt')\n",
    "    ls2file(test_id, 'test-id.txt')\n",
    "    \n",
    "    print('train size: %d, eval size: %d, test size: %d' % (len(train_id), len(eval_id), len(test_id)))\n",
    "    \n",
    "split_dataset()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "avg med for one  198.97109826589596\n"
     ]
    }
   ],
   "source": [
    "import dill\n",
    "# generate ehr graph for gamenet\n",
    "def generate_ehr_graph():\n",
    "    data_multi = pd.read_pickle('data-multi-visit.pkl')\n",
    "    data_single = pd.read_pickle('data-single-visit.pkl')\n",
    "    \n",
    "    rx_voc_size = 0\n",
    "    rx_voc = {}\n",
    "    with open('rx-vocab.txt', 'r') as fin:\n",
    "        for line in fin:\n",
    "            rx_voc[line.rstrip('\\n')] = rx_voc_size\n",
    "            rx_voc_size += 1\n",
    "        \n",
    "    ehr_adj = np.zeros((rx_voc_size, rx_voc_size))\n",
    "    \n",
    "    for idx, row in data_multi.iterrows():\n",
    "        med_set = list(map(lambda x: rx_voc[x], row['ATC4']))\n",
    "        for i, med_i in enumerate(med_set):\n",
    "            for j, med_j in enumerate(med_set):\n",
    "                if j<=i:\n",
    "                    continue\n",
    "                ehr_adj[med_i, med_j] = 1\n",
    "                ehr_adj[med_j, med_i] = 1\n",
    "                \n",
    "    for idx, row in data_single.iterrows():\n",
    "        med_set = list(map(lambda x: rx_voc[x], row['ATC4']))\n",
    "        for i, med_i in enumerate(med_set):\n",
    "            for j, med_j in enumerate(med_set):\n",
    "                if j<=i:\n",
    "                    continue\n",
    "                ehr_adj[med_i, med_j] = 1\n",
    "                ehr_adj[med_j, med_i] = 1\n",
    "    \n",
    "    print('avg med for one ', np.mean(np.sum(ehr_adj,axis=-1)))\n",
    "    \n",
    "    return ehr_adj\n",
    "\n",
    "ehr_adj = generate_ehr_graph()\n",
    "dill.dump(ehr_adj, open('ehr_adj.pkl', 'wb'))  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "845\n"
     ]
    }
   ],
   "source": [
    "# max len medical codes\n",
    "data = data_multi_visit\n",
    "\n",
    "max_len = 0\n",
    "for subject_id in data['SUBJECT_ID'].unique():\n",
    "    item_df = data[data['SUBJECT_ID'] == subject_id]\n",
    "    len_tmp = 0\n",
    "    for index, row in item_df.iterrows():\n",
    "        len_tmp += (len(row['ICD9_CODE']) + len(row['ATC4']))\n",
    "    if len_tmp > max_len:\n",
    "        max_len = len_tmp\n",
    "print(max_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "374\n"
     ]
    }
   ],
   "source": [
    "print(max_len)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "pd.read_pickle(file_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.rename(columns={'NDC':'ATC4'}, inplace=True)\n",
    "data.drop(columns=['PRO_CODE', 'NDC_Len'], axis=1, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(15016, 4)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
